from typing import Union, Dict, Any
from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.bank_risk_service import BankRiskService
from app.schemas.bank_risk_schema import (
    BankRiskRequest, BankRiskResponse, BankRiskMiddleResponse, TaskStatus, BankRiskInputParams, BankTransactionNode
)
from app.tools.redis_client import redis_client
from app.utils.logger import logger

router = APIRouter()
service = BankRiskService()

@router.post("/create", response_model=BankRiskMiddleResponse)
async def create_task(request: BankRiskRequest) -> BankRiskMiddleResponse:
    """
    创建新的银行风险分析任务并存储在Redis中
    
    Args:
        request: 包含任务ID、回调URL和输入参数的请求对象
    """
    try:
        task_response = BankRiskMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.model_dump())
        return task_response
    except Exception as e:
        logger.error(f"创建银行风险分析任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

async def add_task_to_background(
    background_tasks: BackgroundTasks,
    input_params: Dict[str, Any],
    task_id: str
) -> None:
    """将银行风险分析任务添加到后台执行队列"""
    background_tasks.add_task(run_task_in_background, input_params, task_id)

@router.post("/execute/{task_id}", response_model=BankRiskMiddleResponse)
async def execute_task(
    task_id: str,
    background_tasks: BackgroundTasks
) -> BankRiskMiddleResponse:
    """执行指定ID的银行风险分析任务"""
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="银行风险分析任务不存在")

        input_params = task_data.get("input_params")
        if not input_params:
            raise HTTPException(status_code=400, detail="任务输入参数不存在")
            
        result_key = f'{task_id}_result'
        
        # 更新任务状态为执行中
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        # 添加到后台任务队列
        await add_task_to_background(background_tasks, input_params, task_id)
        
        return BankRiskMiddleResponse(**running_status)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行银行风险分析任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/result/{task_id}", response_model=Union[BankRiskResponse, BankRiskMiddleResponse])
async def get_result(task_id: str) -> Union[BankRiskResponse, BankRiskMiddleResponse]:
    """获取银行风险分析任务执行结果"""
    try:
        # 首先获取原始任务数据
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="银行风险分析任务不存在")
            
        # 然后获取结果数据
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="银行风险分析结果不存在")
        
        # 确保 result_data 包含 input_params
        if "input_params" not in result_data:
            result_data["input_params"] = task_data.get("input_params")
        
        # 根据是否存在 output_params 返回不同的响应类型
        if "output_params" in result_data:
            return BankRiskResponse(**result_data)
        else:
            return BankRiskMiddleResponse(**result_data)
            
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取银行风险分析结果失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

async def run_task_in_background(input_params: Dict[str, Any], task_id: str) -> None:
    """在后台执行银行风险分析任务"""
    try:
        # 获取原始任务数据以保留 input_params
        task_data = await redis_client.get_data(task_id)
        
        # 将 dict 转换为 BankRiskInputParams 对象
        input_params_obj = BankRiskInputParams(
            transaction_edges=input_params["transaction_edges"],
            nodes_info=[BankTransactionNode(**node) for node in input_params["nodes_info"]],
            time_window=input_params.get("time_window", 10),
            risk_threshold=input_params.get("risk_threshold", 0.7),
            community_count=input_params.get("community_count", 5),
            alpha=input_params.get("alpha", 0.8)
        )
        
        # 执行分析
        await service.analyze_risk_patterns(
            input_params=input_params_obj,
            task_id=task_id,
        )
        
    except Exception as e:
        logger.error(f"银行风险分析后台任务执行失败: {str(e)}")
        # 更新任务状态为失败，同时保留 input_params
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e),
            "input_params": task_data.get("input_params"),
            "metrics": []
        }
        await redis_client.update_data(f'{task_id}_result', error_status) 